//
// Created by twetec on 25-8-16.
//

#include "../include/network_prober/probe_detect/handle_icmp_probe_detect.h"

#include <sys/socket.h>
#include <arpa/inet.h>
#include <netinet/ip_icmp.h>
#include <netinet/ip6.h>
#include <netinet/icmp6.h>

#include <boost/json.hpp>

using namespace std::literals::chrono_literals;

#include <bs_thread_pool/BS_thread_pool.hpp>

namespace network_prober::probe_detect
{
// 线程池
static BS::thread_pool s_iopool{200};

void tag_invoke(boost::json::value_from_tag, boost::json::value &jv, ICMPProbeDetectResultIPv4 const &c)
{
    auto& ja = jv.emplace_object();

    ja["interface"] = c.interface;
    ja["target"] = c.target;
    ja["delay"] = c.delay;
}

std::ostream& operator<<(std::ostream& out, const ICMPProbeDetectResultIPv4& result)
{
    out
        << "interface: " << result.interface
        << " & uuid: " << result.target
        << " & delay: " << result.delay;

    return out;
}

HandleICMPProbeDetectIPv4::HandleICMPProbeDetectIPv4(const std::string &interface, const std::string &target)
    : m_interface(interface), m_target(target)
{
}

void HandleICMPProbeDetectIPv4::setOnDetectedCallback(std::function<void(ICMPProbeDetectResultIPv4&&)> callback)
{
    m_onDetectedCallback = std::move(callback);
}

/**
 * icmp ipv4 协议头
 */
// typedef struct
// {
//     std::uint8_t type;  // 类型(8=请求，0=应答)
//     std::uint8_t code;      // 代码(通常为0)
//     std::uint16_t checksum; // 校验和
//     std::uint16_t id;       // 标识符
//     std::uint16_t seq;      // 序列号
// } ICMPHeaderIPv4;

typedef struct
{
    struct icmphdr header;  // icmp v4 协议头;  // icmp v4 协议头
    struct timeval timestamp;  // 时间戳
    char payload[56];   //  负载
} ICMPPacketIPv4;

unsigned short computeICMPChecksumIPv4(unsigned short* addr, int len)
{
    int nleft = len;    // 数据包字节长度，用于计数未经过计算的字节数
    int sum = 0;
    unsigned short* w = addr;   // 数据包指针
    unsigned short answer = 0;  // 16位

    // 计算报文中两个字节的和
    while (nleft > 1) {
        sum += *w++;    // 1、从数据包中取16位；2、将取出的16位累加到和中 3、指针再移动16位，指向下两个要累加的数据字节内存
        nleft -= 2; // 将未计算的字节数减少2个，直到只剩1或0个字节
    }

    // 如果报文长度是奇数，最后一个字节作为高8位，再用0填充1个字节(低8位)
    if (nleft == 1) {
        *(unsigned char *)(&answer) = *(unsigned char *)w;  // 把16位类型指针转为8位类型指针，取剩下的1个字节复制给answer
        sum += answer;  // 将低位累加
    }

    sum = (sum >> 16) + (sum & 0xffff); // 高16位 + 低16位 = sum只剩16位
    sum += (sum >> 16);
    answer = ~sum;
    return answer;
}

void HandleICMPProbeDetectIPv4::detect()
{
    // 指向当前对象的shared_ptr
    auto self{shared_from_this()};

    // 发起探测
    s_iopool.submit_task([this, self] {
        // 构建探测结果
        ICMPProbeDetectResultIPv4 result {
            .interface = m_interface,
            .target = m_target,
            .delay = -1
        };

        // 当socket套接字无效时，尝试3次
        for (std::size_t i = 0; i < 3; i ++) {
            // 创建UDP套接字
            int sockfd = socket(AF_INET, SOCK_RAW, IPPROTO_ICMP);
            if (sockfd < 0) {
                // 输出日志
                SPDLOG_INFO("sockfd创建失败: {}", strerror(errno));
                // 继续尝试获取套接字
                continue;
            }

            // 设置接收超时(3s)
            struct timeval tv;
            tv.tv_sec = 3000 / 1000;
            tv.tv_usec = (3000 % 1000) * 1000;
            if (setsockopt(sockfd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv)) < 0) {
               // 关闭socket
               close(sockfd);
               // 输出日志
               SPDLOG_INFO("sockfd: {} - 超时设置时间失败: {}", sockfd, strerror(errno));
               // 判断错误是否是套接字
               if (errno == EBADF) continue;
               // 结束循环
               break;
            }

            // 设定网口地址
            struct sockaddr_in interfaceAddr;
            memset(&interfaceAddr, 0, sizeof(interfaceAddr));
            interfaceAddr.sin_family = AF_INET;
            interfaceAddr.sin_port = htons(0);
            if (inet_pton(AF_INET, m_interface.c_str(), &interfaceAddr.sin_addr) <= 0) {
                // 关闭socket
                close(sockfd);
                // 输出日志
                SPDLOG_INFO("sockfd: {} - 网口地址设置失败: {}", sockfd, strerror(errno));
                // 判断错误是否是套接字
                if (errno == EBADF) continue;
                // 结束循环
                break;
            }

            // 绑定本机网口
            if (bind(sockfd, (sockaddr *)&interfaceAddr, sizeof(interfaceAddr)) < 0) {
                // 关闭socket
                close(sockfd);
                // 输出日志
                SPDLOG_INFO("sockfd: {} - 网口地址绑定失败: {}", sockfd, strerror(errno));
                // 判断错误是否是套接字
                if (errno == EBADF) continue;
                // 结束循环
                break;
            }

            // 设置目标地址
            struct sockaddr_in targetAddr;
            memset(&targetAddr, 0, sizeof(targetAddr));
            targetAddr.sin_family = AF_INET;
            if (inet_pton(AF_INET, m_target.c_str(), &targetAddr.sin_addr) <= 0) {
                // 关闭socket
                close(sockfd);
                // 输出日志
                SPDLOG_INFO("sockfd: {} - 目标地址设置失败: {}", sockfd, strerror(errno));
                // 判断错误是否是套接字
                if (errno == EBADF) continue;
                // 结束循环
                break;
            }

            // 填充ICMP包
            ICMPPacketIPv4 sendPacket;
            memset(&sendPacket, 0, sizeof(sendPacket));
            sendPacket.header.type = ICMP_ECHO;
            sendPacket.header.code = 0;
            sendPacket.header.un.echo.id = getpid() & 0xFFFF;
            strncpy(sendPacket.payload, "PING", sizeof(sendPacket.payload)-1);
            // 准备校验和
            sendPacket.header.un.echo.sequence = 0;
            gettimeofday(&sendPacket.timestamp, NULL);
            sendPacket.header.checksum = 0;
            sendPacket.header.checksum = computeICMPChecksumIPv4((unsigned short*)&sendPacket, sizeof(sendPacket));

            // 请求计时 - 开始时间
            std::chrono::time_point<std::chrono::steady_clock> detectStartTime = std::chrono::steady_clock::now();

            // 发送ICMP包
            if (sendto(sockfd, &sendPacket, sizeof(sendPacket), 0, (struct sockaddr *)&targetAddr, sizeof(targetAddr)) <= 0) {
                // 关闭socket
                close(sockfd);
                // 输出日志
                SPDLOG_INFO("sockfd: {} - ICMP数据包发送失败: {}", sockfd, strerror(errno));
                // 判断错误是否是套接字
                if (errno == EBADF) continue;
                // 结束循环
                break;
            }

            // 尝试接收响应
            // ip头部 + icmp头部 + icmp数据
            char recvBuf[1024];
            socklen_t targetAddrLen = sizeof(targetAddr);

            // 接收响应
            ssize_t recvdBytesCount = recvfrom(sockfd, recvBuf, sizeof(recvBuf), 0, (struct sockaddr *)&targetAddr, &targetAddrLen);

            // 关闭socket
            close(sockfd);

            // 如果未接收到数据
            if (recvdBytesCount <= 0) {
                // 输出日志
                // SPDLOG_INFO("sockfd: {} - ICMP数据包接收失败: {}", sockfd, strerror(errno));
                // 判断错误是否是套接字
                if (errno == EBADF) continue;
                // 结束执行
                break;
            }

            // 如果接收到数据，则作处理
            // 解析IP头获取ICMP包位置
            struct iphdr *ipHeader = (struct iphdr *)recvBuf;
            size_t ipHeaderLen = ipHeader->ihl * 4;
            // 如果包长度不够，则说明接收失败
            if (recvdBytesCount < ipHeaderLen + ICMP_MINLEN) {
                // 输出日志
                SPDLOG_INFO("sockfd: {} - ICMP数据包长度不够: {}", sockfd, strerror(errno));
                // 结束循环
                break;
            }

            // 提取icmp协议头 (跳过ip头)
            struct icmphdr *icmpHdr = (struct icmphdr *)(recvBuf + ipHeaderLen);
            // 检查是否是我们的响应
            if (icmpHdr->type == ICMP_ECHOREPLY && icmpHdr->un.echo.id == sendPacket.header.un.echo.id) {
                // 计算延迟
                result.delay = std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::steady_clock::now() - detectStartTime).count();
                // 结束循环
                break;
            }
        }

        m_onDetectedCallback(std::move(result));
    });
}

HandleICMPProbeDetectIPv4::~HandleICMPProbeDetectIPv4()
{
}

std::ostream& operator<<(std::ostream& out, const ICMPProbeDetectResultIPv6& result)
{
    out
        << "interface: " << result.interface
        << " & uuid: " << result.target
        << " & delay: " << result.delay;

    return out;
}

void tag_invoke(boost::json::value_from_tag, boost::json::value &jv, ICMPProbeDetectResultIPv6 const &c)
{
    auto& ja = jv.emplace_object();

    ja["interface"] = c.interface;
    ja["target"] = c.target;
    ja["delay"] = c.delay;
}

HandleICMPProbeDetectIPv6::HandleICMPProbeDetectIPv6(const std::string &interface, const std::string &target)
    : m_interface(interface), m_target(target)
{
}

void HandleICMPProbeDetectIPv6::setOnDetectedCallback(std::function<void(ICMPProbeDetectResultIPv6&&)> callback)
{
    m_onDetectedCallback = std::move(callback);
}

typedef struct
{
    struct icmp6_hdr header;  // icmp v4 协议头
    struct timeval tv;  // 时间戳
    char payload[64-sizeof(icmp6_hdr)];   //  负载
} ICMPPacketIPv6;

void HandleICMPProbeDetectIPv6::detect()
{
    // 指向当前对象的shared_ptr
    auto self{shared_from_this()};

    // 发起探测
    s_iopool.submit_task([this, self] {
        // 构建探测结果
        ICMPProbeDetectResultIPv6 result {
            .interface = m_interface,
            .target = m_target,
            .delay = -1
        };

        // 当socket套接字无效时，尝试3次
        for (std::size_t i = 0; i < 3; i ++) {
            // 创建UDP套接字
            int sockfd = socket(AF_INET6, SOCK_RAW, IPPROTO_ICMPV6);
            if (sockfd < 0) {
                // 输出日志
                SPDLOG_INFO("sockfd创建失败: {}", strerror(errno));
                // 继续尝试获取套接字
                continue;
            }

            // 设置接收超时(3s)
            struct timeval tv;
            tv.tv_sec = 3000 / 1000;
            tv.tv_usec = (3000 % 1000) * 1000;
            if (setsockopt(sockfd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv)) < 0) {
               // 关闭socket
               close(sockfd);
               // 输出日志
               SPDLOG_INFO("sockfd: {} - 超时设置时间失败: {}", sockfd, strerror(errno));
               // 判断错误是否是套接字
               if (errno == EBADF) continue;
               // 结束循环
               break;
            }

            // 设定网口地址
            struct sockaddr_in6 interfaceAddr;
            memset(&interfaceAddr, 0, sizeof(interfaceAddr));
            interfaceAddr.sin6_family = AF_INET6;
            interfaceAddr.sin6_port = htons(0);
            if (inet_pton(AF_INET6, m_interface.c_str(), &interfaceAddr.sin6_addr) <= 0) {
                // 关闭socket
                close(sockfd);
                // 输出日志
                SPDLOG_INFO("sockfd: {} - 网口地址设置失败: {}", sockfd, strerror(errno));
                // 判断错误是否是套接字
                if (errno == EBADF) continue;
                // 结束循环
                break;
            }

            // 绑定本机网口
            if (bind(sockfd, (sockaddr *)&interfaceAddr, sizeof(interfaceAddr)) < 0) {
                // 关闭socket
                close(sockfd);
                // 输出日志
                SPDLOG_INFO("sockfd: {} - 网口地址绑定失败: {}", sockfd, strerror(errno));
                // 判断错误是否是套接字
                if (errno == EBADF) continue;
                // 结束循环
                break;
            }

            // 设置目标地址
            struct sockaddr_in6 targetAddr;
            memset(&targetAddr, 0, sizeof(targetAddr));
            targetAddr.sin6_family = AF_INET6;
            if (inet_pton(AF_INET6, m_target.c_str(), &targetAddr.sin6_addr) <= 0) {
                // 关闭socket
                close(sockfd);
                // 输出日志
                SPDLOG_INFO("sockfd: {} - 目标地址设置失败: {}", sockfd, strerror(errno));
                // 判断错误是否是套接字
                if (errno == EBADF) continue;
                // 结束循环
                break;
            }

            // 填充ICMP包
            ICMPPacketIPv6 sendPacket;
            memset(&sendPacket, 0, sizeof(sendPacket));
            sendPacket.header.icmp6_type = ICMP6_ECHO_REQUEST;
            sendPacket.header.icmp6_code = 0;
            sendPacket.header.icmp6_cksum = 0;  // checksum will be calculated later
            sendPacket.header.icmp6_id = htons(getpid() & 0xFFFF);  // Identifier
            sendPacket.header.icmp6_seq = htons(1);  // Sequence number
            strncpy(sendPacket.payload, "PING", sizeof(sendPacket.payload)-1);

            // 请求计时 - 开始时间
            std::chrono::time_point<std::chrono::steady_clock> detectStartTime = std::chrono::steady_clock::now();

            // 发送ICMP包
            if (sendto(sockfd, &sendPacket, sizeof(sendPacket), 0, (struct sockaddr *)&targetAddr, sizeof(targetAddr)) <= 0) {
                // 关闭socket
                close(sockfd);
                // 输出日志
                SPDLOG_INFO("sockfd: {} - ICMP数据包发送失败: {}", sockfd, strerror(errno));
                // 判断错误是否是套接字
                if (errno == EBADF) continue;
                // 结束循环
                break;
            }

            // 尝试接收响应
            // ip头部 + icmp头部 + icmp数据
            char recvBuf[1024];
            socklen_t targetAddrLen = sizeof(targetAddr);

            // 接收响应
            ssize_t recvdBytesCount = recvfrom(sockfd, recvBuf, sizeof(recvBuf), 0, (struct sockaddr *)&targetAddr, &targetAddrLen);

            // 关闭socket
            close(sockfd);

            // 如果未接收到数据
            if (recvdBytesCount <= 0) {
                // 输出日志
                // SPDLOG_INFO("sockfd: {} - 数据包接收失败: {}", sockfd, strerror(errno));
                // 判断错误是否是套接字
                if (errno == EBADF) continue;
                // 结束执行
                break;
            }

            // 如果接收到数据，则作处理
            // 解析IP头获取ICMP包位置
            struct icmp6_hdr *icmpHdr = (struct icmp6_hdr *)recvBuf;
            // 检查是否是我们的响应
            if (icmpHdr->icmp6_type == ICMP6_ECHO_REPLY) {
                // 计算延迟
                result.delay = std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::steady_clock::now() - detectStartTime).count();
                // 结束循环
                break;
            }
        }

        m_onDetectedCallback(std::move(result));
    });
}


HandleICMPProbeDetectIPv6::~HandleICMPProbeDetectIPv6()
{

}

}